package app

import (
	"sync"

	"github.com/go-rod/rod"
	"github.com/go-rod/rod/lib/proto"
)

// WebService web服务.
type WebService struct {
	pages map[string]*rod.Page
	lock  sync.RWMutex
}

// NewWebService 新web服务.
func NewWebService() *WebService {
	return &WebService{
		pages: make(map[string]*rod.Page),
	}
}

// GetPage 网页.
func (s *WebService) GetPage(url string) *rod.Page {
	s.lock.Lock()
	defer s.lock.Unlock()

	if _, has := s.pages[url]; !has {
		page := rod.New().MustConnect().MustPage(url)
		page.MustWaitLoad()
		s.pages[url] = page
	}

	return s.pages[url]
}

// Screenshot 截图.
func (s *WebService) Screenshot(url string, selector string) []byte {
	page := s.GetPage(url)
	if selector == "" || !page.MustHas(selector) {
		return page.MustScreenshot("")
	}

	e := page.MustElement(selector)
	bs, _ := e.Screenshot(proto.PageCaptureScreenshotFormatPng, 10)

	return bs
}
